from abc import ABC, abstractmethod
from typing import Dict, List, Iterator, Optional
from datetime import datetime
import pandas as pd
from decimal import Decimal
import ccxt
from functools import wraps
import time

def retry_on_exception(max_retries=3, delay=1):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if i < max_retries - 1:
                        print(f"获取数据时出错: {str(e)}, 正在重试...")
                        time.sleep(delay)
                    else:
                        raise e
        return wrapper
    return decorator

class DataHandler(ABC):
    """数据处理器基类"""
    
    def __init__(self):
        """初始化数据处理器"""
        self.continue_backtest = True
        self._data = {}  # 交易所 -> 交易对 -> DataFrame
        self._latest_data = {}  # 交易所 -> 交易对 -> 最新数据
        self._current_idx = {}  # 交易所 -> 交易对 -> 当前索引
        
    @property
    def data(self) -> Dict[str, Dict[str, pd.DataFrame]]:
        """获取数据"""
        return self._data
        
    @property
    def latest_data(self) -> Dict[str, Dict[str, List[Dict]]]:
        """获取最新数据"""
        return self._latest_data
        
    @property
    def current_idx(self) -> Dict[str, Dict[str, int]]:
        """获取当前索引"""
        return self._current_idx
        
    @abstractmethod
    def update_bars(self):
        """更新最新的K线数据"""
        raise NotImplementedError("Should implement update_bars()")
        
    @abstractmethod
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的一条数据"""
        raise NotImplementedError("Should implement get_latest_bar()")
        
    @abstractmethod
    def get_latest_bars(self, exchange: str, symbol: str, *, n: int = 50) -> List[Dict]:
        """获取最新的n条数据"""
        raise NotImplementedError("Should implement get_latest_bars()")
        
    @abstractmethod
    def get_all_bars(self) -> Dict[str, Dict[str, pd.DataFrame]]:
        """获取所有K线数据"""
        raise NotImplementedError("Should implement get_all_bars()")

class HistoricalDataHandler(DataHandler):
    """历史数据处理器"""
    
    def __init__(
        self,
        exchanges: Dict[str, ccxt.Exchange],
        symbols: List[str],
        timeframe: str,
        start_date: datetime,
        end_date: datetime
    ):
        """初始化历史数据处理器"""
        super().__init__()
        
        self.exchanges = exchanges
        self.symbols = symbols
        self.timeframe = timeframe
        self.start_date = start_date
        self.end_date = end_date
        
        # 设置当前交易所和交易对
        self.current_exchange = list(exchanges.keys())[0]
        self.current_symbol = symbols[0]
        
        # 加载历史数据
        self._load_data()
        
    def _load_data(self) -> None:
        """加载历史数据"""
        print("\n加载历史数据...")
        
        for exchange_name, exchange in self.exchanges.items():
            print(f"处理交易所: {exchange_name}")
            
            # 确保交易所键存在
            if exchange_name not in self._data:
                self._data[exchange_name] = {}
                self._latest_data[exchange_name] = {}
                self._current_idx[exchange_name] = {}
            
            for symbol in self.symbols:
                try:
                    print(f"正在从{exchange_name}获取{symbol}的历史数据...")
                    
                    # 获取历史数据
                    ohlcv = exchange.fetch_ohlcv(
                        symbol=symbol,
                        timeframe=self.timeframe,
                        since=int(self.start_date.timestamp() * 1000),
                        limit=None
                    )
                    
                    # 转换为DataFrame
                    df = pd.DataFrame(
                        ohlcv,
                        columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                    )
                    
                    # 转换时间戳
                    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                    df.set_index('timestamp', inplace=True)
                    
                    # 过滤时间范围
                    df = df[
                        (df.index >= pd.Timestamp(self.start_date)) &
                        (df.index <= pd.Timestamp(self.end_date))
                    ]
                    
                    print(f"获取到{len(df)}条{symbol}的历史数据")
                    
                    # 存储数据，使用正确的键结构
                    self._data[exchange_name][symbol] = df
                    self._latest_data[exchange_name][symbol] = []
                    self._current_idx[exchange_name][symbol] = 0
                    
                except Exception as e:
                    print(f"加载{exchange_name}的{symbol}数据时出错: {str(e)}")
                    raise
                    
    def update_bars(self):
        """更新最新的K线数据"""
        for exchange_name in self.exchanges:
            for symbol in self.symbols:
                try:
                    if exchange_name not in self._data or symbol not in self._data[exchange_name]:
                        continue
                        
                    df = self._data[exchange_name][symbol]
                    current_idx = self._current_idx[exchange_name][symbol]
                    
                    if current_idx >= len(df):
                        self.continue_backtest = False
                        return
                        
                    latest_bar = df.iloc[current_idx].to_dict()
                    latest_bar['timestamp'] = df.index[current_idx]
                    self._latest_data[exchange_name][symbol].append(latest_bar)
                    self._current_idx[exchange_name][symbol] += 1
                    
                except Exception as e:
                    print(f"更新{exchange_name}的{symbol}数据时出错: {str(e)}")
                    raise
                    
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的一条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-1] if latest_data else None
        except (KeyError, IndexError):
            return None
            
    def get_latest_bars(self, exchange: str, symbol: str, *, n: int = 50) -> List[Dict]:
        """获取最新的n条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-n:] if len(latest_data) >= n else latest_data
        except (KeyError, IndexError):
            return []
            
    def get_all_bars(self) -> Dict[str, Dict[str, pd.DataFrame]]:
        """获取所有K线数据"""
        return self._data

class LiveDataHandler(DataHandler):
    """实时数据处理器"""
    
    def __init__(self, exchanges: Dict[str, ccxt.Exchange], symbols: List[str], timeframe: str = '1h'):
        """
        初始化实时数据处理器
        
        Args:
            exchanges: 交易所实例字典
            symbols: 交易对列表
            timeframe: 时间周期
        """
        super().__init__()
        self.exchanges = exchanges
        self.symbols = symbols
        self.timeframe = timeframe
        
        # 初始化数据
        self._init_data()
        
    def _init_data(self):
        """初始化数据"""
        for exchange_name in self.exchanges:
            if exchange_name not in self._data:
                self._data[exchange_name] = {}
                self._latest_data[exchange_name] = {}
                self._current_idx[exchange_name] = {}
            
            for symbol in self.symbols:
                self._data[exchange_name][symbol] = pd.DataFrame()
                self._latest_data[exchange_name][symbol] = []
                self._current_idx[exchange_name][symbol] = 0
                
    @retry_on_exception()
    def _fetch_ohlcv(self, exchange: ccxt.Exchange, symbol: str, **kwargs) -> List:
        """获取K线数据"""
        timeframe = kwargs.pop('timeframe', self.timeframe)
        return exchange.fetch_ohlcv(symbol, timeframe=timeframe, **kwargs)
        
    def update_bars(self):
        """更新最新的K线数据"""
        for exchange_name, exchange in self.exchanges.items():
            for symbol in self.symbols:
                try:
                    # 获取最新数据
                    ohlcv = self._fetch_ohlcv(
                        exchange,
                        symbol,
                        limit=1
                    )
                    
                    if not ohlcv:
                        continue
                        
                    # 转换为DataFrame
                    df = pd.DataFrame(
                        ohlcv,
                        columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                    )
                    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                    df.set_index('timestamp', inplace=True)
                    
                    # 更新数据
                    self._data[exchange_name][symbol] = pd.concat([
                        self._data[exchange_name][symbol],
                        df
                    ]).drop_duplicates()
                    
                    # 添加到最新数据列表
                    latest_bar = df.iloc[-1].to_dict()
                    latest_bar['timestamp'] = df.index[-1]
                    self._latest_data[exchange_name][symbol].append(latest_bar)
                    
                except Exception as e:
                    print(f"更新{exchange_name}的{symbol}数据时出错: {str(e)}")
                    continue
                    
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的一条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-1] if latest_data else None
        except (KeyError, IndexError):
            return None
            
    def get_latest_bars(self, exchange: str, symbol: str, *, n: int = 50) -> List[Dict]:
        """获取最新的n条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-n:] if len(latest_data) >= n else latest_data
        except (KeyError, IndexError):
            return []
            
    def get_all_bars(self) -> Dict[str, Dict[str, pd.DataFrame]]:
        """获取所有K线数据"""
        return self._data
